{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Exemplar Simulations\n",
    "Exemplar simulations investigating the performance degradation of Memristive Deep Neural Networks (MDNNs) when non-ideal device characteristics are accounted for using the CIFAR-10 dataset are provided below. Results can be plotted using `plot_all_exemplar.m`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 1. Define the MobileNetV2 Network Architecture"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "\r\n",
    "\r\n",
    "class Block(nn.Module):\r\n",
    "    def __init__(self, in_planes, out_planes, expansion, stride):\r\n",
    "        super(Block, self).__init__()\r\n",
    "        self.stride = stride\r\n",
    "        planes = expansion * in_planes\r\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, stride=1, padding=0)\r\n",
    "        self.bn1 = nn.BatchNorm2d(planes, affine=False)\r\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, groups=planes)\r\n",
    "        self.bn2 = nn.BatchNorm2d(planes, affine=False)\r\n",
    "        self.conv3 = nn.Conv2d(planes, out_planes, kernel_size=1, stride=1, padding=0)\r\n",
    "        self.bn3 = nn.BatchNorm2d(out_planes, affine=False)\r\n",
    "\r\n",
    "        self.shortcut = nn.Sequential()\r\n",
    "        if stride == 1 and in_planes != out_planes:\r\n",
    "            self.shortcut = nn.Sequential(\r\n",
    "                nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, padding=0),\r\n",
    "                nn.BatchNorm2d(out_planes, affine=False),\r\n",
    "            )\r\n",
    "\r\n",
    "    def forward(self, x):\r\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\r\n",
    "        out = F.relu(self.bn2(self.conv2(out)))\r\n",
    "        out = self.bn3(self.conv3(out))\r\n",
    "        out = out + self.shortcut(x) if self.stride==1 else out\r\n",
    "        return out\r\n",
    "\r\n",
    "\r\n",
    "class MobileNetV2(nn.Module):\r\n",
    "    cfg = [(1,  16, 1, 1),\r\n",
    "           (6,  24, 2, 1),\r\n",
    "           (6,  32, 3, 2),\r\n",
    "           (6,  64, 4, 2),\r\n",
    "           (6,  96, 3, 1),\r\n",
    "           (6, 160, 3, 2),\r\n",
    "           (6, 320, 1, 1)]\r\n",
    "\r\n",
    "    def __init__(self, num_classes=10):\r\n",
    "        super(MobileNetV2, self).__init__()\r\n",
    "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)\r\n",
    "        self.bn1 = nn.BatchNorm2d(32, affine=False)\r\n",
    "        self.layers = self._make_layers(in_planes=32)\r\n",
    "        self.conv2 = nn.Conv2d(320, 1280, kernel_size=1, stride=1, padding=0)\r\n",
    "        self.bn2 = nn.BatchNorm2d(1280, affine=False)\r\n",
    "        self.linear = nn.Linear(1280, num_classes)\r\n",
    "\r\n",
    "    def _make_layers(self, in_planes):\r\n",
    "        layers = []\r\n",
    "        for expansion, out_planes, num_blocks, stride in self.cfg:\r\n",
    "            strides = [stride] + [1]*(num_blocks-1)\r\n",
    "            for stride in strides:\r\n",
    "                layers.append(Block(in_planes, out_planes, expansion, stride))\r\n",
    "                in_planes = out_planes\r\n",
    "        return nn.Sequential(*layers)\r\n",
    "\r\n",
    "    def forward(self, x):\r\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\r\n",
    "        out = self.layers(out)\r\n",
    "        out = F.relu(self.bn2(self.conv2(out)))\r\n",
    "        out = F.avg_pool2d(out, 4)\r\n",
    "        out = out.view(out.size(0), -1)\r\n",
    "        out = self.linear(out)\r\n",
    "        return out"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 2. Train MobileNetV2 Using CIFAR-10"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "from torch.autograd import Variable\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "import torch.optim as optim\r\n",
    "import numpy as np\r\n",
    "import random\r\n",
    "import torchvision\r\n",
    "import torchvision.transforms as transforms\r\n",
    "\r\n",
    "\r\n",
    "def set_all_seeds(seed):\r\n",
    "  random.seed(seed)\r\n",
    "  np.random.seed(seed)\r\n",
    "  torch.manual_seed(seed)\r\n",
    "  torch.cuda.manual_seed(seed)\r\n",
    "  torch.backends.cudnn.deterministic = True\r\n",
    "\r\n",
    "def test(model, test_loader):\r\n",
    "    correct = 0\r\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\r\n",
    "        output = model(data.to(device))\r\n",
    "        pred = output.data.max(1)[1]\r\n",
    "        correct += pred.eq(target.to(device).data.view_as(pred)).cpu().sum()\r\n",
    "\r\n",
    "    return 100. * float(correct) / float(len(test_loader.dataset))\r\n",
    "\r\n",
    "set_all_seeds(0)\r\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\r\n",
    "epochs = 100\r\n",
    "transform_train = transforms.Compose([\r\n",
    "    transforms.RandomCrop(32, padding=4),\r\n",
    "    transforms.RandomHorizontalFlip(),\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "transform_test = transforms.Compose([\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)\r\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, shuffle=True, num_workers=1)\r\n",
    "test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)\r\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False, num_workers=1)\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "criterion = nn.CrossEntropyLoss()\r\n",
    "learning_rate = 0.1\r\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)\r\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1)\r\n",
    "best_accuracy = 0\r\n",
    "for epoch in range(0, epochs):\r\n",
    "    print('Epoch: [%d]\\t\\t' % (epoch + 1), end='')\r\n",
    "    model.train()\r\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\r\n",
    "        optimizer.zero_grad()\r\n",
    "        output = model(data.to(device))\r\n",
    "        loss = criterion(output, target.to(device))\r\n",
    "        loss.backward()\r\n",
    "        optimizer.step()\r\n",
    "\r\n",
    "    scheduler.step()\r\n",
    "    model.eval()\r\n",
    "    accuracy = test(model, test_loader)\r\n",
    "    print('%2.2f%%' % accuracy)\r\n",
    "    if accuracy > best_accuracy:\r\n",
    "        print('Saving model...')\r\n",
    "        torch.save(model.state_dict(), 'trained_model.pt')\r\n",
    "        best_accuracy = accuracy"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 3. Figure 1 [A,E]"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "import torch.optim as optim\r\n",
    "from torchvision import datasets, transforms\r\n",
    "import memtorch\r\n",
    "from memtorch.utils import LoadCIFAR10\r\n",
    "import numpy as np\r\n",
    "import pandas as pd\r\n",
    "import torchvision\r\n",
    "import copy\r\n",
    "from memtorch.mn.Module import patch_model\r\n",
    "from memtorch.map.Parameter import naive_map\r\n",
    "from memtorch.bh.crossbar.Program import naive_program\r\n",
    "from memtorch.bh.nonideality.NonIdeality import apply_nonidealities\r\n",
    "\r\n",
    "\r\n",
    "def test(model, test_loader):\r\n",
    "    correct = 0\r\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\r\n",
    "        output = model(data.to(device))\r\n",
    "        pred = output.data.max(1)[1]\r\n",
    "        correct += pred.eq(target.to(device).data.view_as(pred)).cpu().sum()\r\n",
    "\r\n",
    "    return 100. * float(correct) / float(len(test_loader.dataset))\r\n",
    "\r\n",
    "device = torch.device('cuda')\r\n",
    "transform_test = transforms.Compose([\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)\r\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False, num_workers=1)\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "try:\r\n",
    "    model.load_state_dict(torch.load('trained_model.pt'), strict=False)\r\n",
    "    model.eval()\r\n",
    "except:\r\n",
    "    raise Exception('trained_model.pt has not been found.')\r\n",
    "\r\n",
    "print('Test Set Accuracy: \\t%2.2f%%' % test(model, test_loader))\r\n",
    "\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "model.load_state_dict(torch.load('trained_model.pt'), strict=True)\r\n",
    "model.eval()\r\n",
    "print(test(model, test_loader))\r\n",
    "r_on = 1.4e4\r\n",
    "r_off = 5e7\r\n",
    "\r\n",
    "def trial(r_on, r_off, tile_shape, ADC_resolution, sigma):\r\n",
    "    model_ = copy.deepcopy(model)\r\n",
    "    reference_memristor = memtorch.bh.memristor.VTEAM\r\n",
    "    if sigma == 0.:\r\n",
    "        reference_memristor_params = {'time_series_resolution': 1e-10, 'r_off': r_off, 'r_on': r_on}\r\n",
    "    else:\r\n",
    "        reference_memristor_params = {'time_series_resolution': 1e-10,\r\n",
    "                                  'r_off': memtorch.bh.StochasticParameter(loc=r_off, scale=sigma*2, min=1),\r\n",
    "                                  'r_on': memtorch.bh.StochasticParameter(loc=r_on, scale=sigma, min=1)}\r\n",
    "\r\n",
    "    patched_model = patch_model(copy.deepcopy(model_),\r\n",
    "                              memristor_model=reference_memristor,\r\n",
    "                              memristor_model_params=reference_memristor_params,\r\n",
    "                              module_parameters_to_patch=[torch.nn.Linear, torch.nn.Conv2d],\r\n",
    "                              mapping_routine=naive_map,\r\n",
    "                              transistor=True,\r\n",
    "                              programming_routine=None,\r\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn,\r\n",
    "                              tile_shape=tile_shape,\r\n",
    "                              max_input_voltage=0.3,\r\n",
    "                              ADC_resolution=int(ADC_resolution),\r\n",
    "                              ADC_overflow_rate=0.,\r\n",
    "                              quant_method='linear')\r\n",
    "    \r\n",
    "    patched_model.tune_()\r\n",
    "    return test(patched_model, test_loader)\r\n",
    "\r\n",
    "df = pd.DataFrame(columns=['tile_shape', 'ADC_resolution', 'sigma', 'test_set_accuracy'])\r\n",
    "tile_shape = [(256, 64)]\r\n",
    "ADC_resolution = np.linspace(2, 10, num=5, endpoint=True, dtype=int)\r\n",
    "sigma = np.logspace(6, 7, endpoint=True, num=5)\r\n",
    "for tile_shape_ in tile_shape:\r\n",
    "    for ADC_resolution_ in ADC_resolution:\r\n",
    "        for sigma_ in sigma:\r\n",
    "            print('tile_shape: %s; ADC_resolution: %d; sigma: %d' % (tile_shape_, ADC_resolution_, sigma_))\r\n",
    "            df = df.append({'tile_shape': tile_shape_, \r\n",
    "                            'ADC_resolution': ADC_resolution_, \r\n",
    "                            'sigma': sigma_, \r\n",
    "                            'test_set_accuracy': trial(r_on, r_off, tile_shape_, ADC_resolution_, sigma_)}, ignore_index=True)\r\n",
    "            df.to_csv('1_AE.csv', index=False)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 4. Figure 1 [B,F]"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "import torch.optim as optim\r\n",
    "from torchvision import datasets, transforms\r\n",
    "import memtorch\r\n",
    "from memtorch.utils import LoadCIFAR10\r\n",
    "import numpy as np\r\n",
    "import pandas as pd\r\n",
    "import torchvision\r\n",
    "import copy\r\n",
    "from memtorch.mn.Module import patch_model\r\n",
    "from memtorch.map.Parameter import naive_map\r\n",
    "from memtorch.bh.crossbar.Program import naive_program\r\n",
    "from memtorch.bh.nonideality.NonIdeality import apply_nonidealities\r\n",
    "\r\n",
    "\r\n",
    "def test(model, test_loader):\r\n",
    "    correct = 0\r\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\r\n",
    "        output = model(data.to(device))\r\n",
    "        pred = output.data.max(1)[1]\r\n",
    "        correct += pred.eq(target.to(device).data.view_as(pred)).cpu().sum()\r\n",
    "\r\n",
    "    return 100. * float(correct) / float(len(test_loader.dataset))\r\n",
    "\r\n",
    "device = torch.device('cuda')\r\n",
    "transform_test = transforms.Compose([\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)\r\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False, num_workers=1)\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "try:\r\n",
    "    model.load_state_dict(torch.load('trained_model.pt'), strict=False)\r\n",
    "    model.eval()\r\n",
    "except:\r\n",
    "    raise Exception('trained_model.pt has not been found.')\r\n",
    "\r\n",
    "print('Test Set Accuracy: \\t%2.2f%%' % test(model, test_loader))\r\n",
    "\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "model.load_state_dict(torch.load('trained_model.pt'), strict=True)\r\n",
    "model.eval()\r\n",
    "print(test(model, test_loader))\r\n",
    "r_on = 1.4e4\r\n",
    "r_off = 5e7\r\n",
    "\r\n",
    "def trial(r_on, r_off, tile_shape, ADC_resolution, conductance_states):\r\n",
    "    model_ = copy.deepcopy(model)\r\n",
    "    reference_memristor = memtorch.bh.memristor.VTEAM\r\n",
    "    reference_memristor_params = {'time_series_resolution': 1e-10, 'r_off': r_off, 'r_on': r_on}\r\n",
    "    patched_model = patch_model(copy.deepcopy(model_),\r\n",
    "                              memristor_model=reference_memristor,\r\n",
    "                              memristor_model_params=reference_memristor_params,\r\n",
    "                              module_parameters_to_patch=[torch.nn.Linear, torch.nn.Conv2d],\r\n",
    "                              mapping_routine=naive_map,\r\n",
    "                              transistor=True,\r\n",
    "                              programming_routine=None,\r\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn,\r\n",
    "                              tile_shape=tile_shape,\r\n",
    "                              max_input_voltage=0.3,\r\n",
    "                              ADC_resolution=int(ADC_resolution),\r\n",
    "                              ADC_overflow_rate=0.,\r\n",
    "                              quant_method='linear')\r\n",
    "\r\n",
    "    patched_model = apply_nonidealities(patched_model,\r\n",
    "                                        non_idealities=[memtorch.bh.nonideality.NonIdeality.FiniteConductanceStates],\r\n",
    "                                        conductance_states = int(conductance_states))\r\n",
    "    \r\n",
    "    patched_model.tune_()\r\n",
    "    return test(patched_model, test_loader)\r\n",
    "\r\n",
    "df = pd.DataFrame(columns=['tile_shape', 'ADC_resolution', 'conductance_states', 'test_set_accuracy'])\r\n",
    "torch.backends.cudnn.benchmark = False\r\n",
    "tile_shape = [(128, 128), (256, 64)]\r\n",
    "ADC_resolution = np.linspace(2, 10, num=5, endpoint=True, dtype=int)\r\n",
    "conductance_states = np.linspace(2, 10, num=5, endpoint=True, dtype=int)\r\n",
    "for tile_shape_ in tile_shape:\r\n",
    "    for ADC_resolution_ in ADC_resolution:\r\n",
    "        for conductance_states_ in conductance_states:\r\n",
    "            print('tile_shape: %s; ADC_resolution: %d; conductance_states: %d' % (tile_shape_, ADC_resolution_, conductance_states_))\r\n",
    "            df = df.append({'tile_shape': tile_shape_, \r\n",
    "                            'ADC_resolution': ADC_resolution_, \r\n",
    "                            'conductance_states': conductance_states_, \r\n",
    "                            'test_set_accuracy': trial(r_on, r_off, tile_shape_, ADC_resolution_, conductance_states_)}, ignore_index=True)\r\n",
    "            df.to_csv('1_BF.csv', index=False)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 5. Figure 1 [C,G]"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "import torch.optim as optim\r\n",
    "from torchvision import datasets, transforms\r\n",
    "import memtorch\r\n",
    "from memtorch.utils import LoadCIFAR10\r\n",
    "import numpy as np\r\n",
    "import pandas as pd\r\n",
    "import torchvision\r\n",
    "import copy\r\n",
    "from memtorch.mn.Module import patch_model\r\n",
    "from memtorch.map.Parameter import naive_map\r\n",
    "from memtorch.bh.crossbar.Program import naive_program\r\n",
    "from memtorch.bh.nonideality.NonIdeality import apply_nonidealities\r\n",
    "\r\n",
    "\r\n",
    "def test(model, test_loader):\r\n",
    "    correct = 0\r\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\r\n",
    "        output = model(data.to(device))\r\n",
    "        pred = output.data.max(1)[1]\r\n",
    "        correct += pred.eq(target.to(device).data.view_as(pred)).cpu().sum()\r\n",
    "\r\n",
    "    return 100. * float(correct) / float(len(test_loader.dataset))\r\n",
    "\r\n",
    "device = torch.device('cuda')\r\n",
    "transform_test = transforms.Compose([\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)\r\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False, num_workers=1)\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "try:\r\n",
    "    model.load_state_dict(torch.load('trained_model.pt'), strict=False)\r\n",
    "    model.eval()\r\n",
    "except:\r\n",
    "    raise Exception('trained_model.pt has not been found.')\r\n",
    "\r\n",
    "print('Test Set Accuracy: \\t%2.2f%%' % test(model, test_loader))\r\n",
    "\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "model.load_state_dict(torch.load('trained_model.pt'), strict=True)\r\n",
    "model.eval()\r\n",
    "print(test(model, test_loader))\r\n",
    "r_on = 1.4e4\r\n",
    "r_off = 5e7\r\n",
    "\r\n",
    "def trial(r_on, r_off, tile_shape, ADC_resolution, failure_percentage):\r\n",
    "    model_ = copy.deepcopy(model)\r\n",
    "    reference_memristor = memtorch.bh.memristor.VTEAM\r\n",
    "    reference_memristor_params = {'time_series_resolution': 1e-10, 'r_off': r_off, 'r_on': r_on}\r\n",
    "    patched_model = patch_model(copy.deepcopy(model_),\r\n",
    "                              memristor_model=reference_memristor,\r\n",
    "                              memristor_model_params=reference_memristor_params,\r\n",
    "                              module_parameters_to_patch=[torch.nn.Linear, torch.nn.Conv2d],\r\n",
    "                              mapping_routine=naive_map,\r\n",
    "                              transistor=True,\r\n",
    "                              programming_routine=None,\r\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn,\r\n",
    "                              tile_shape=tile_shape,\r\n",
    "                              max_input_voltage=0.3,\r\n",
    "                              ADC_resolution=int(ADC_resolution),\r\n",
    "                              ADC_overflow_rate=0.,\r\n",
    "                              quant_method='linear')\r\n",
    "\r\n",
    "    patched_model = apply_nonidealities(patched_model,\r\n",
    "                                        non_idealities=[memtorch.bh.nonideality.NonIdeality.DeviceFaults],\r\n",
    "                                        lrs_proportion=failure_percentage,\r\n",
    "                                        hrs_proportion=0.,\r\n",
    "                                        electroform_proportion=0.)\r\n",
    "    \r\n",
    "    patched_model.tune_()\r\n",
    "    return test(patched_model, test_loader)\r\n",
    "\r\n",
    "df = pd.DataFrame(columns=['tile_shape', 'ADC_resolution', 'failure_percentage', 'test_set_accuracy'])\r\n",
    "tile_shape = [(128, 128), (256, 64)]\r\n",
    "ADC_resolution = np.linspace(2, 10, num=5, endpoint=True, dtype=int)\r\n",
    "failure_percentage = np.linspace(0, 0.25, 5)\r\n",
    "for tile_shape_ in tile_shape:\r\n",
    "    for ADC_resolution_ in ADC_resolution:\r\n",
    "        for failure_percentage_ in failure_percentage:\r\n",
    "            print('tile_shape: %s; ADC_resolution: %d; failure_percentage: %d' % (tile_shape_, ADC_resolution_, failure_percentage_))\r\n",
    "            df = df.append({'tile_shape': tile_shape_, \r\n",
    "                            'ADC_resolution': ADC_resolution_, \r\n",
    "                            'failure_percentage': failure_percentage_, \r\n",
    "                            'test_set_accuracy': trial(r_on, r_off, tile_shape_, ADC_resolution_, failure_percentage_)}, ignore_index=True)\r\n",
    "            df.to_csv('1_CG.csv', index=False)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 6. Figure 1 [D,H]"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import torch\r\n",
    "import torch.nn as nn\r\n",
    "import torch.nn.functional as F\r\n",
    "import torch.optim as optim\r\n",
    "from torchvision import datasets, transforms\r\n",
    "import memtorch\r\n",
    "from memtorch.utils import LoadCIFAR10\r\n",
    "import numpy as np\r\n",
    "import pandas as pd\r\n",
    "import torchvision\r\n",
    "import copy\r\n",
    "from memtorch.mn.Module import patch_model\r\n",
    "from memtorch.map.Parameter import naive_map\r\n",
    "from memtorch.bh.crossbar.Program import naive_program\r\n",
    "from memtorch.bh.nonideality.NonIdeality import apply_nonidealities\r\n",
    "\r\n",
    "\r\n",
    "def test(model, test_loader):\r\n",
    "    correct = 0\r\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\r\n",
    "        output = model(data.to(device))\r\n",
    "        pred = output.data.max(1)[1]\r\n",
    "        correct += pred.eq(target.to(device).data.view_as(pred)).cpu().sum()\r\n",
    "\r\n",
    "    return 100. * float(correct) / float(len(test_loader.dataset))\r\n",
    "\r\n",
    "device = torch.device('cuda')\r\n",
    "transform_test = transforms.Compose([\r\n",
    "    transforms.ToTensor(),\r\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),\r\n",
    "])\r\n",
    "test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)\r\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False, num_workers=1)\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "try:\r\n",
    "    model.load_state_dict(torch.load('trained_model.pt'), strict=False)\r\n",
    "    model.eval()\r\n",
    "except:\r\n",
    "    raise Exception('trained_model.pt has not been found.')\r\n",
    "\r\n",
    "print('Test Set Accuracy: \\t%2.2f%%' % test(model, test_loader))\r\n",
    "\r\n",
    "model = MobileNetV2().to(device)\r\n",
    "model.load_state_dict(torch.load('trained_model.pt'), strict=True)\r\n",
    "model.eval()\r\n",
    "print(test(model, test_loader))\r\n",
    "r_on = 1.4e4\r\n",
    "r_off = 5e7\r\n",
    "\r\n",
    "def trial(r_on, r_off, tile_shape, ADC_resolution, failure_percentage):\r\n",
    "    model_ = copy.deepcopy(model)\r\n",
    "    reference_memristor = memtorch.bh.memristor.VTEAM\r\n",
    "    reference_memristor_params = {'time_series_resolution': 1e-10, 'r_off': r_off, 'r_on': r_on}\r\n",
    "    patched_model = patch_model(copy.deepcopy(model_),\r\n",
    "                              memristor_model=reference_memristor,\r\n",
    "                              memristor_model_params=reference_memristor_params,\r\n",
    "                              module_parameters_to_patch=[torch.nn.Linear, torch.nn.Conv2d],\r\n",
    "                              mapping_routine=naive_map,\r\n",
    "                              transistor=True,\r\n",
    "                              programming_routine=None,\r\n",
    "                              scheme=memtorch.bh.Scheme.DoubleColumn,\r\n",
    "                              tile_shape=tile_shape,\r\n",
    "                              max_input_voltage=0.3,\r\n",
    "                              ADC_resolution=int(ADC_resolution),\r\n",
    "                              ADC_overflow_rate=0.,\r\n",
    "                              quant_method='linear')\r\n",
    "\r\n",
    "    patched_model = apply_nonidealities(patched_model,\r\n",
    "                                        non_idealities=[memtorch.bh.nonideality.NonIdeality.DeviceFaults],\r\n",
    "                                        lrs_proportion=0.,\r\n",
    "                                        hrs_proportion=failure_percentage,\r\n",
    "                                        electroform_proportion=0.)\r\n",
    "    \r\n",
    "    patched_model.tune_()\r\n",
    "    return test(patched_model, test_loader)\r\n",
    "\r\n",
    "df = pd.DataFrame(columns=['tile_shape', 'ADC_resolution', 'failure_percentage', 'test_set_accuracy'])\r\n",
    "tile_shape = [(128, 128), (256, 64)]\r\n",
    "ADC_resolution = np.linspace(2, 10, num=5, endpoint=True, dtype=int)\r\n",
    "failure_percentage = np.linspace(0, 0.25, 5)\r\n",
    "for tile_shape_ in tile_shape:\r\n",
    "    for ADC_resolution_ in ADC_resolution:\r\n",
    "        for failure_percentage_ in failure_percentage:\r\n",
    "            print('tile_shape: %s; ADC_resolution: %d; failure_percentage: %d' % (tile_shape_, ADC_resolution_, failure_percentage_))\r\n",
    "            df = df.append({'tile_shape': tile_shape_, \r\n",
    "                            'ADC_resolution': ADC_resolution_, \r\n",
    "                            'failure_percentage': failure_percentage_, \r\n",
    "                            'test_set_accuracy': trial(r_on, r_off, tile_shape_, ADC_resolution_, failure_percentage_)}, ignore_index=True)\r\n",
    "            df.to_csv('1_DH.csv', index=False)"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}